# $Id: tree.py 7557b5d4d82b 2012/12/02 09:07:16 pgurumur $
# Copyright (c) 2009 Prabhu Gurumurthy  <pgurumur@gmail.com>
# 
# Permission to use, copy, modify, and distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
#

__version__ = "$Revision: 7557b5d4d82b $"
__author__ = "$Author: pgurumur $"

import re

from djanet.lib import exception
from djanet.lib import ip

ASNList = []
IPv4List = []

class ASNTree(object):
   def __init__(self, ASNumber, Server, Info):
      self._err = exception.TreeError()
      self._asn = 0
      self._start = 0
      self._end = 0
      self._server = None
      self._info = None
      self._query = None

      if ASNumber:
         asregex = re.compile(r'(\d{1,10})-(\d{1,10})')
         match = asregex.search(ASNumber)
         if match:
            (start, end) = match.groups()
            self._start = int(start)
            self._end = int(end)
         else:
            self._asn = int(ASNumber)

      if Server:
         self._server = Server.strip()

      if Info:
         self._info = Info

      if self._server == 'whois.arin.net':
         self._query = "a "
      else:
         self._query = "as"

   def __getitem__(self, Key):
      retval = None
      if Key in self:
         retval = self._server

      return retval

   def __contains__(self, ASNumber):
      retval = False
      value = 0
      try:
         value = int(ASNumber)
      except ValueError, message:
         value = 0

      if value:
         if self._asn:
            if self._asn == value:
               retval = True
         else:
            try:
               if value in range(self._start, (self._end + 1)):
                  retval = True
            except OverflowError, message:
               if (value >= self._start) and (value <= self._end):
                  retval = True

      if not self._server:
         retval = False

      return retval

   def __str__(self):
      retval = None
      if self._asn:
         retval = "ASN: %d, info: %s" %(self._asn, self._info)
      else:
         retval = "ASN from %d to %d, info: %s" %(
               self._start, self._end, self._info)

      return retval

   def __repr__(self):
      return repr(str(self))

   def __unicode__(self):
      return unicode(str(self))

   query = property(fget = lambda self: self._query)

class IPv4Tree(object):
   def __init__(self, Network, Server):
      self._err = exception.TreeError()
      self._net = None
      self._server = 'whois.arin.net'
      self._query = None

      if Network:
         ipnet = None
         ipregex = re.compile(r'^(\d{1,3})/(\d{1,2})')
         match = ipregex.search(Network)
         if match:
            (net, mask) = match.groups()
            self._net = ip.ip("%d/%d" %(int(net), int(mask)))

      if Server:
         self._server = Server.strip()

      if self._server == 'whois.arin.net':
         self._query = "n "

   def __contains__(self, Network):
      retval = False
      value = None
      try:
         value = ip.ip(Network)
      except ip.IPError, message:
         pass

      if value:
         if value in self._net:
            retval = True

      return retval

   def __str__(self):
      retval = None
      if self._net and self._server:
         retval = "<%s: %s>" %(str(self._net), self._server)
      elif self._net:
         retval = str(self._net)

      return retval

   def __repr__(self):
      return repr(str(self))

   def __unicode__(self):
      return unicode(str(self))

   def __getitem__(self, Key):
      retval = None
      if Key in self:
         retval = self._server

      return retval

   query = property(fget = lambda self: self._query)
